import os
from openai import OpenAI  # 注意：这里用的是OpenAI兼容接口，实际调用的是阿里云服务
from tqdm import tqdm


def text_embeddings(texts):
    """
    将文本转换为向量（向量化）

    参数:
        texts: 文本列表（可以是单个字符串或多个字符串组成的列表）

    返回:
        向量列表，每个向量对应输入的文本
    """
    # 初始化客户端（连接阿里云通义千问服务）
    client = OpenAI(
        # 从环境变量获取API密钥，也可以直接写成：api_key="你的密钥"
        api_key=os.getenv("DASHSCOPE_API_KEY"),
        # 阿里云通义千问的兼容接口地址
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
    )

    # 如果输入是单个字符串，转换成列表形式
    if isinstance(texts, str):
        texts = [texts]

    # 通义千问text-embedding-v4模型一次最多处理10个文本，所以要分批处理
    batch_size = 10
    embeddings = []  # 存储所有向量

    # 分批处理文本（显示进度条）
    for i in tqdm(range(0, len(texts), batch_size), desc="生成向量中"):
        # 截取当前批次的文本
        batch = texts[i:i + batch_size]

        # 调用模型生成向量
        completion = client.embeddings.create(
            model="text-embedding-v4",  # 指定使用的模型
            input=batch,  # 输入文本
            dimensions=1024,  # 向量维度（1024是比较常用的尺寸）
            encoding_format="float"  # 向量格式为浮点数
        )

        # 提取当前批次的向量并添加到结果列表
        batch_embeddings = [item.embedding for item in completion.data]
        embeddings.extend(batch_embeddings)

    return embeddings


# 测试代码
if __name__ == "__main__":
    # 假设这是上一步切分得到的文本片段（保险相关内容）
    sample_chunks = [
        {"page_num": 1, "chunk": "平安保险成立于1988年，总部位于深圳。"},
        {"page_num": 1, "chunk": "旗下有多个保险产品，包括寿险、车险、健康险等。"},
        {"page_num": 2, "chunk": "平安福是公司的明星寿险产品，2024年推出了新版。"}
    ]

    # 提取文本内容（只需要chunk部分）
    texts = [chunk["chunk"] for chunk in sample_chunks]

    # 调用向量化函数
    vectors = text_embeddings(texts)

    # 打印结果
    print(f"生成了{len(vectors)}个向量")
    # 打印第一个向量的前5个数值（向量很长，只看一部分）
    if vectors:
        print("第一个向量的前5个值：", vectors[0][:5])
        print("向量长度（维度）：", len(vectors[0]))  # 应该是1024
